package com.ena.luademo.aspect;

import cn.hutool.core.util.StrUtil;
import com.ena.luademo.annotation.Limiter;
import com.ena.luademo.enums.LimitTypeEnum;
import com.google.common.collect.ImmutableList;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.io.Serializable;
import java.lang.reflect.Method;

/**
 * @Author zgr
 * @create 2023/8/8 12:14
 */
@Aspect
@Configuration
@Slf4j
public class LimitAspect {

    private static final String UNKNOWN = "unknown";

    private final RedisTemplate<String, Serializable> redisTemplate;


    public LimitAspect(RedisTemplate<String, Serializable> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    @Pointcut("execution(public * *(..)) && @annotation(com.ena.luademo.annotation.Limiter)")
    public void methodPointCut() {}

    @Around(value = "methodPointCut() && @annotation(limiter)")
    public Object around(ProceedingJoinPoint pjg, Limiter limiter) {
        Object obj = null;
        MethodSignature signature = (MethodSignature) pjg.getSignature();
        Method method = signature.getMethod();

        LimitTypeEnum limitType = limiter.limitType();

        String key="";
        //根据限流类型获取不同的key
        switch (limitType) {
            case IP:
                key = getIpAddress();
                break;
            case CUSTOMER:
                key = limiter.key();
                break;
            default:
                key = StrUtil.swapCase(method.getName());
        }

        ImmutableList<String> keys = ImmutableList.of(StrUtil.join(limiter.prefix(), key));
        try {
            String luaScript = buildLuaScript();
            RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);
            Long count = redisTemplate.execute(redisScript, keys, limiter.count(), limiter.period());
            log.info(" count is {} for name={} and key={}", count, limiter.name(), key);
            if (count != null && count <= limiter.count()) {
                obj = pjg.proceed();
            } else {
                throw new RuntimeException("over error");
            }
        } catch (Throwable e) {
            if (e instanceof RuntimeException) {
                throw new RuntimeException(e.getLocalizedMessage());
            }
            throw new RuntimeException("server exception");
        }
        return obj;
    }

    /**
     *  编写 redis Lua 限流脚本
     */
    public String buildLuaScript() {
        StringBuilder lua = new StringBuilder();
        lua.append("local c");
        lua.append("\nc = redis.call('get',KEYS[1])");
        // 调用不超过最大值，则直接返回
        lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        lua.append("\nreturn c;");
        lua.append("\nend");
        // 执行计算器自加
        lua.append("\nc = redis.call('incr',KEYS[1])");
        lua.append("\nif tonumber(c) == 1 then");
        // 从第一次调用开始限流，设置对应键值的过期
        lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        lua.append("\nend");
        lua.append("\nreturn c;");
        return lua.toString();
    }


    public String getIpAddress() {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        String ip = request.getHeader("x-forwarded-for");
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getHeader("WL-Proxy-Client-IP");
        }
        if (ip == null || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
            ip = request.getRemoteAddr();
        }
        return ip;
    }


}
